import os
from PySide6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QPushButton, QFileDialog,
    QTextEdit, QLabel, QLineEdit, QHBoxLayout
)

class FileTool(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文件扫描与恢复工具")
        self.resize(800, 600)

        # 默认屏蔽类型
        self.ignore_extensions = [".py", ".dll", ".lib"]

        layout = QVBoxLayout()

        # 屏蔽类型设置
        layout.addWidget(QLabel("屏蔽的文件类型（用逗号分隔，如 .py,.dll,.lib）:"))
        self.ignore_input = QLineEdit(",".join(self.ignore_extensions))
        layout.addWidget(self.ignore_input)

        # 扫描目录
        layout.addWidget(QLabel("扫描目录:"))
        scan_layout = QHBoxLayout()
        self.scan_dir_input = QLineEdit()
        scan_btn = QPushButton("选择目录")
        scan_btn.clicked.connect(self.choose_scan_dir)
        scan_layout.addWidget(self.scan_dir_input)
        scan_layout.addWidget(scan_btn)
        layout.addLayout(scan_layout)

        self.scan_text = QTextEdit()
        layout.addWidget(self.scan_text)

        scan_generate_btn = QPushButton("扫描并生成文本")
        scan_generate_btn.clicked.connect(self.scan_and_generate)
        layout.addWidget(scan_generate_btn)

        # 分割线
        layout.addWidget(QLabel("—— 反向生成文件 ——"))

        # 反向生成部分
        layout.addWidget(QLabel("选择文本文件:"))
        restore_layout = QHBoxLayout()
        self.restore_file_input = QLineEdit()
        restore_btn = QPushButton("选择文件")
        restore_btn.clicked.connect(self.choose_restore_file)
        restore_layout.addWidget(self.restore_file_input)
        restore_layout.addWidget(restore_btn)
        layout.addLayout(restore_layout)

        output_layout = QHBoxLayout()
        output_layout.addWidget(QLabel("输出目录:"))
        self.output_dir_input = QLineEdit("recovered_files")
        output_layout.addWidget(self.output_dir_input)
        layout.addLayout(output_layout)

        restore_generate_btn = QPushButton("从文本生成文件")
        restore_generate_btn.clicked.connect(self.restore_files)
        layout.addWidget(restore_generate_btn)

        self.setLayout(layout)

    # ===== 扫描目录 =====
    def choose_scan_dir(self):
        folder = QFileDialog.getExistingDirectory(self, "选择目录")
        if folder:
            self.scan_dir_input.setText(folder)

    def scan_and_generate(self):
        root_dir = self.scan_dir_input.text() or "."
        # 动态更新屏蔽类型
        self.ignore_extensions = [ext.strip() for ext in self.ignore_input.text().split(",") if ext.strip()]

        result_text = []

        for dirpath, dirnames, filenames in os.walk(root_dir):
            filenames.sort()
            for filename in filenames:
                if any(filename.endswith(ext) for ext in self.ignore_extensions):
                    continue

                file_path = os.path.join(dirpath, filename)
                relative_path = os.path.relpath(file_path, root_dir)
                # 使用唯一边界标记
                result_text.append(f"=====FILE_START: {relative_path}=====")
                try:
                    with open(file_path, "r", encoding="utf-8") as f:
                        result_text.append(f.read())
                except Exception as e:
                    result_text.append(f"// 无法读取文件: {e}")
                result_text.append(f"=====FILE_END=====\n")

        final_text = "\n".join(result_text)
        self.scan_text.setPlainText(final_text)

        with open("all_files_combined.txt", "w", encoding="utf-8") as f:
            f.write(final_text)

    # ===== 反向生成文件 =====
    def choose_restore_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择文本文件", filter="Text Files (*.txt);;All Files (*)")
        if file_path:
            self.restore_file_input.setText(file_path)

    def restore_files(self):
        input_file = self.restore_file_input.text()
        output_root = self.output_dir_input.text() or "recovered_files"
        if not os.path.exists(output_root):
            os.makedirs(output_root)

        current_file = None
        content_lines = []

        with open(input_file, "r", encoding="utf-8") as f:
            for line in f:
                line_strip = line.strip()
                if line_strip.startswith("=====FILE_START:"):
                    if current_file:
                        self.write_file(current_file, content_lines)
                        content_lines = []
                    relative_path = line_strip[len("=====FILE_START:"):].replace("=====", "").strip()
                    current_file = os.path.join(output_root, relative_path)
                elif line_strip == "=====FILE_END=====":
                    if current_file:
                        self.write_file(current_file, content_lines)
                        current_file = None
                        content_lines = []
                else:
                    if current_file is not None:
                        content_lines.append(line)

        self.scan_text.setPlainText(f"已恢复文件到 {output_root}")

    def write_file(self, file_path, lines):
        os.makedirs(os.path.dirname(file_path), exist_ok=True)
        with open(file_path, "w", encoding="utf-8") as f:
            f.writelines(lines)


if __name__ == "__main__":
    app = QApplication([])
    window = FileTool()
    window.show()
    app.exec()
